package blue.mqtt4j.internal.client.impl;

import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.handler.codec.mqtt.MqttMessage;
import io.netty.handler.codec.mqtt.MqttMessageType;
import io.netty.handler.codec.mqtt.MqttQoS;
import io.netty.util.ReferenceCountUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.NavigableSet;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.ConcurrentSkipListSet;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * @author Jin Zheng
 * @since 1.0 2020-10-21
 */
public class PublishHandler implements Runnable
{
	private static Logger logger = LoggerFactory.getLogger(PublishHandler.class);

	private Channel channel;

	private volatile boolean stop = false;
	private final Thread thread;
	private final int delay;
	private final ConcurrentMap<Integer, ChannelFutureListener> listenerMap = new ConcurrentHashMap<>();
	private final BlockingQueue<PublishMessage> publishQueue = new LinkedBlockingQueue<>();
	private final NavigableSet<Integer> pubAckSet = new ConcurrentSkipListSet<>();
	private final NavigableSet<Integer> pubRecSet = new ConcurrentSkipListSet<>();
	private final NavigableSet<Integer> pubCompSet = new ConcurrentSkipListSet<>();

	public PublishHandler(int delay)
	{
		this.delay = delay;
		this.thread = new Thread(this);
		this.thread.start();
	}

	public void publish(MqttMessage mqttMessage, ChannelFutureListener listener)
	{
		if (stop)
		{
			logger.error("Publisher has been stopped");
			return;
		}

		PublishMessage message = new PublishMessage(mqttMessage, delay);
		if (channel == null)
		{
			publishQueue.add(message);
			if (listener != null)
			{
				listenerMap.put(message.messageId(), listener);
			}
			logger.warn("Channel was disconnected, add to publishQueue, type: {}, messageId: {}",
					message.type(), message.messageId());
			return;
		}

		if (message.type() == MqttMessageType.PUBREL ||
				message.qos() == MqttQoS.AT_LEAST_ONCE || message.qos() == MqttQoS.EXACTLY_ONCE)
		{
			this.addQueue(message);
		}
		ChannelFuture future = channel.writeAndFlush(mqttMessage)
				.addListener(PublisherListener.createWithQoS0(message, publishQueue));
		if (listener != null)
		{
			future.addListener(listener);
		}
	}

	private void addQueue(PublishMessage message)
	{
		ReferenceCountUtil.retain(message.message());
		publishQueue.offer(message);
		logger.debug("Add publishQueue, messageId: {}", message.messageId());
	}

	@Override
	public void run()
	{
		while (!stop)
		{
			try
			{
				PublishMessage message = publishQueue.take();
				this.sleep(message);
				switch (message.type())
				{
					case PUBREL:
						this.publishWithPubRel(message);
						break;
					case PUBLISH:
						switch (message.qos())
						{
							case AT_MOST_ONCE:
								this.publishWithQoS0(message);
								break;
							case AT_LEAST_ONCE:
								this.publishWithQoS1(message);
								break;
							case EXACTLY_ONCE:
								this.publishWithQoS2(message);
								break;
							default:
								logger.warn("Unknown MqttQoS: {}", message.qos());
								break;
						}
						break;
					default:
						break;
				}
			}
			catch (InterruptedException e)
			{
				logger.warn("Interrupted from BlockingQueue/Thread.sleep()");
			}
		}
	}

	private void addChannelFutureListener(PublishMessage message, ChannelFuture future)
	{
		ChannelFutureListener listener = listenerMap.get(message.messageId());
		if (listener != null)
		{
			future.addListener(listener);
			listenerMap.remove(message.messageId());
		}
	}

	private void publishWithQoS0(PublishMessage message)
	{
		ChannelFuture future = channel.writeAndFlush(message.message())
				.addListener(PublisherListener.createWithQoS0(message, publishQueue));
		this.addChannelFutureListener(message, future);
	}

	private void publishWithQoS1(PublishMessage message)
	{
		if (pubAckSet.contains(message.messageId()))
		{
			ReferenceCountUtil.release(message.message());
			pubAckSet.remove(message.messageId());
			return;
		}

		this.addQueue(message.buildNew());
		ChannelFuture future = channel.writeAndFlush(ReferenceCountUtil.retain(message.message()))
				.addListener(PublisherListener.createWithQoS1(message, pubAckSet));
		this.addChannelFutureListener(message, future);
	}

	private void publishWithQoS2(PublishMessage message)
	{
		if (pubRecSet.contains(message.messageId()))
		{
			ReferenceCountUtil.release(message.message());
			pubRecSet.remove(message.messageId());
			return;
		}

		this.addQueue(message.buildNew());
		ChannelFuture future = channel.writeAndFlush(ReferenceCountUtil.retain(message.message()))
				.addListener(PublisherListener.createWithQoS2(message, pubRecSet));
		this.addChannelFutureListener(message, future);
	}

	private void publishWithPubRel(PublishMessage message)
	{
		if (pubCompSet.contains(message.messageId()))
		{
			ReferenceCountUtil.release(message.message());
			pubCompSet.remove(message.messageId());
			return;
		}

		this.addQueue(message.buildNew());
		channel.writeAndFlush(message.message())
				.addListener(PublisherListener.createWithPubRel(message, pubCompSet));
	}

	private void sleep(PublishMessage message) throws InterruptedException
	{
		long millis = message.sendTime() - System.currentTimeMillis();
		if (millis <= 0)
			return;

		Thread.sleep(millis);
	}

	public void setPubAck(Integer messageId)
	{
		pubAckSet.add(messageId);
	}

	public void setPubRec(Integer messageId)
	{
		pubRecSet.add(messageId);
	}

	public void setPubComp(Integer messageId)
	{
		pubCompSet.add(messageId);
	}

	public void stop()
	{
		stop = true;
		thread.interrupt();
	}

	public void setChannel(Channel channel)
	{
		this.channel = channel;
	}
}
